#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <getopt.h>
#include <memory.h>
#include <unistd.h>
#include "frodoII.h"

static uint8_t	wbuf[BLOCK_SIZE];

static void usage(){
#ifdef TARGET_PLATFORM
	fprintf(stderr, "Usage: frodoII -K|-F|-X <image>\n");
#else
	fprintf(stderr, "Usage: frodoII <parameters>\n\n");
	fprintf(stderr, "Extract files from image to stdout:\n");
	fprintf(stderr, "	-K <image>		- Extract the kernel from the image\n");
	fprintf(stderr, "	-F <image>		- Extract the RootFS from the image\n");
	fprintf(stderr, "Check image:\n");
	fprintf(stderr, "	-X <image>		- Check image integrity\n");
	fprintf(stderr, "Calculate checksum:\n");
	fprintf(stderr, "	-C <image>		- Calculate a file checksum\n");
	fprintf(stderr, "Create frodoII image file:\n");
	fprintf(stderr, "	-i <kernel>		- The kernel file to include in the image\n");
	fprintf(stderr, "	-f <rootfs>		- The RootFS file to include in the image\n");
	fprintf(stderr, "	-s <settings>		- The Settings file to include in the image\n");
	fprintf(stderr, "	-o <bin>		- The output file\n");
	fprintf(stderr, "	-p <productid>		- The product id to set in the image\n");
	fprintf(stderr, "	-c <customid>		- The custom id to set in the image\n");
	fprintf(stderr, "	-m <modelid>		- The model id to set in the image\n");
	fprintf(stderr, "	-u <subid> 		- The sub id to set in the image (optional, default 1)\n");
	fprintf(stderr, "	-n <newversion>   	- The version id to set in the image (optional, default 4)\n");
	fprintf(stderr, "	-h			- Displays this text\n\n");
#endif
}

static int copy_file(int fdo, const char *inFileName){
	int fdi;
	int i;
	int len;
	ssize_t rc;

	fdi = open(inFileName, O_RDONLY);
	if(fdi == -1){
		fprintf(stderr, "Error open file:%s", inFileName);
		return 0;
	}
	
	i = 0;
	while(1)
	{
		len = BLOCK_SIZE;
		len = read(fdi, wbuf, len);
		if(len < 0){
			close(fdi);
			return 0;
		}
		if(len == 0)
			break;
		rc = write(fdo, wbuf, len);
		if(rc != len ){
			close(fdi);
			return 0;
		}
		if(len != BLOCK_SIZE)
			break;
		i += len;
	}
	close(fdi);
	return 1;
}

#ifndef TARGET_PLATFORM
static uint32_t checksum(const char *file_name)
{
  int		fd;
  uint32_t	csum;
  uint32_t	len, i, j;
  uint32_t	*buf = (uint32_t*)wbuf;
		
  fd = open(file_name, O_RDONLY);
  if(fd < 0){
      fprintf(stderr, "Error: Can't open %s file\n", file_name);
      return 0;
  }

  csum = 0;
  i = 0;
  while(1){
      len = read(fd, buf, BLOCK_SIZE);

      for(j=0; j<len/4; j++){
	  csum ^= buf[j];
      }
      if(len != BLOCK_SIZE)
	  break;
      i += len;
  }
  close(fd);
	
  return csum;
}
#endif

static uint32_t checksum2(FILE *fi, long img_offs, int img_len){

  uint32_t	csum;
  uint32_t	len, i, j;
  uint32_t	*buf = (uint32_t*)wbuf;
  int lenf, lenp;

  fseek(fi, img_offs, SEEK_SET);
		
  lenf = img_len/BLOCK_SIZE;
  lenp = img_len%BLOCK_SIZE;
  csum = 0;
  i = 0;

  for(i=0; i<lenf; i++){
  	len = fread(buf, 1, BLOCK_SIZE, fi);
  	for(j=0; j<len/4; j++){
		csum ^= buf[j];
  	}
  }
  if(lenp){
  	len = fread(buf, 1, BLOCK_SIZE, fi);
  	for(j=0; j<lenp/4; j++){
		csum ^= buf[j];
  	}
  }
  return csum;
}

static int do_extract(const char *file, int element){
	uint32_t offs, len;
	frodo2_control_header_t hdr;
	FILE* fi;
	int i, rc;

	if(file == NULL){
		usage();
		return 1;
	}

	fi = fopen(file, "r");
	if(fi == NULL){
		fprintf(stderr, "Error: Can't open %s file\n", file);
		return 1;
	}
	rc = fread(&hdr, sizeof(hdr), 1, fi);
	if(rc != 1){
		fclose(fi);
		fprintf(stderr, "Invalid image header\n");
		return 1;
	}

	switch(element){
	case 0:
		offs = hdr.kernel_offs;
		len = hdr.kernel_len;
		break;
	case 1:
		offs = hdr.rootfs_offs;
		len = hdr.rootfs_len;
		break;
	default:
		fclose(fi);
		return 1;
	}
	fseek(fi, offs, SEEK_SET);	
	for(i=0; i<len; i++){
		fputc(fgetc(fi), stdout);
	}
	fclose(fi);
	return 0;
}

// return 0 - image ok
// 	1 - bad image
static int do_check(const char *file){
	frodo2_control_header_t hdr;
	struct stat st;
	int rc;
	FILE* fi;
	char *err_string;

	if(file == NULL){
		usage();
		return 1;
	}
	fi = fopen(file, "r");
	if(fi == NULL){
		fprintf(stderr, "Error: Can't open %s file\n", file);
		return 1;
	}
// 1. Check header
	err_string = "Invalid image header";
	rc = fread(&hdr, sizeof(hdr), 1, fi);
	if(rc != 1){
		goto failed; // poor man exception :)
	}
	if(memcmp(FRODOII_MAGIC, hdr.magic_num, sizeof(FRODOII_MAGIC))!=0){
		goto failed;
	}
// TODO: get from command line and check for modelid, productid etc.
	rc = lstat(file, &st);
	if(rc < 0)
		goto failed;
	if((st.st_size < (sizeof(hdr)+hdr.kernel_len+hdr.rootfs_len+hdr.settings_len))||
	   (st.st_size < (hdr.kernel_len+hdr.kernel_offs))||
	   (st.st_size < (hdr.rootfs_len+hdr.rootfs_offs))||
	   (st.st_size < (hdr.settings_len+hdr.settings_offs))
	)
		goto failed;
		
// 2. Check kernel checksum
	err_string = "Invalid kernel checksum";
	if(hdr.kernel_offs != 0 && hdr.kernel_len != 0){
		if( hdr.kernel_csum != checksum2(fi, hdr.kernel_offs, hdr.kernel_len))
			goto failed;
	}
// 3. Check rootfs checksum
	err_string = "Invalid rootfs checksum";
	if(hdr.rootfs_offs != 0 && hdr.rootfs_len != 0){
		if( hdr.rootfs_csum != checksum2(fi, hdr.rootfs_offs, hdr.rootfs_len))
			goto failed;
	}
// 4. Check settings checksum
	err_string = "Invalid settings checksum";
	if(hdr.settings_offs != 0 && hdr.settings_len != 0){
		if( hdr.settings_csum != checksum2(fi, hdr.settings_offs, hdr.settings_len))
			goto failed;
	}

	fclose(fi);
	fprintf(stderr, "Ok\n");
	return 0;

failed:
	fclose(fi);
	fprintf(stderr, "%s\n", err_string);
	return 1;
}

#ifndef TARGET_PLATFORM
static int do_calc_cs(const char *file){
	uint32_t cs;
	if(file == NULL){
		usage();
		return 1;
	}
	cs = checksum(file);
	printf("Checksum: %x\n", cs);
	return 0;
}
#endif

int main(int argc, char **argv){
        int c;

#ifndef TARGET_PLATFORM
	char *kernel = NULL;
	char *rootfs = NULL;
	char *bin = NULL;
	char *settings = NULL;
	frodo2_control_header_t hdr;
	struct stat st;
	int rc;
	int fdo;

	memset(&hdr, 0, sizeof(frodo2_control_header_t));
	hdr.sub_id = 1;
	hdr.new_version = 4;
	memcpy(hdr.magic_num, FRODOII_MAGIC, sizeof(FRODOII_MAGIC));

	while ((c = getopt(argc, argv, "K:F:X:C:i:f:s:o:p:c:m:u:n:h")) != -1) {
		switch (c) {
	
			case 'K':
				return do_extract(optarg, 0);
				break;				
			case 'F':
				return do_extract(optarg, 1);
				break;				
			case 'X':
				return do_check(optarg);
				break;			
			case 'C':
				return do_calc_cs(optarg);
				break;			
			case 'i':
				kernel = optarg;
				break;
			case 'f':
				rootfs = optarg;
				break;
			case 's':
				settings = optarg;
				break;
			case 'o':
				bin = optarg;
				break;
			case 'p':
				hdr.product_id = atoi(optarg);
				break;
			case 'c':
				hdr.custom_id = atoi(optarg);
				break;
			case 'm':
				hdr.model_id = atoi(optarg);
				break;
			case 'u':
				hdr.sub_id = atoi(optarg);
				break;
			case 'n':
				hdr.new_version = atoi(optarg);
				break;
			case 'h':
			default:
				usage();
				return 1;
		}
	}	
	if(rootfs == NULL || kernel == NULL || bin == NULL){
		usage();
		return 1;
	}
	//print id to check
  	printf("Product ID = %x\n", hdr.product_id);
  	printf("Custom  ID = %x\n", hdr.custom_id);
  	printf("Model   ID = %x\n", hdr.model_id);
  	printf("Sub     ID = %x\n", hdr.sub_id);
	
	hdr.kernel_offs = sizeof(hdr);
	rc = lstat(kernel, &st);
	if(rc < 0){
	      fprintf(stderr, "Kernel file size error\n");
	      return 1;
	}
	hdr.kernel_len = st.st_size;
	if(st.st_size > MAX_KERNEL_LEN ){
	      fprintf(stderr, "Kernel size must be less than %d bytes\n", MAX_KERNEL_LEN);
	      return 1;
	}

	hdr.rootfs_offs = sizeof(hdr) + hdr.kernel_len;
	rc = lstat(rootfs, &st);
	if(rc < 0){
	      fprintf(stderr, "Rootfs file size error\n");
	      return 1;
	}
	hdr.rootfs_len = st.st_size;
	if(st.st_size > MAX_ROOTFS_LEN ){
	      fprintf(stderr, "Rootfs size must be less than %d bytes\n", MAX_ROOTFS_LEN);
	      return 1;
	}

	if(settings == NULL){
		hdr.settings_offs = 0;
  		hdr.settings_len = 0;
	}else{
		hdr.settings_offs = sizeof(hdr) + hdr.kernel_len + hdr.rootfs_len;
  		rc = lstat(settings, &st);
		if(rc < 0){
		      fprintf(stderr, "Settings file size error\n");
		      return 1;
		}
  		hdr.settings_len = st.st_size;
 	}
	hdr.next_offs = 0;

	//calculate check sums
	hdr.kernel_csum = checksum(kernel);
	hdr.rootfs_csum = checksum(rootfs);
	if(settings == NULL)
		hdr.settings_csum = 0;
	else
		hdr.settings_csum = checksum(settings);

	//combine files
	fdo = open(bin, O_RDWR | O_CREAT | O_TRUNC, 0664);
  	if( fdo == -1 ){
      		fprintf(stderr, "Error open file %s\n", bin);
      		return 1;
    	}
	rc = write(fdo, &hdr, sizeof(hdr));
	copy_file(fdo, kernel);
	copy_file(fdo, rootfs);
	if(settings != NULL)
		copy_file(fdo, settings);
	close(fdo);
#else
	while ((c = getopt(argc, argv, "K:F:X:C:i:f:s:o:p:c:m:u:n:h")) != -1) {
		switch (c) {
	
			case 'K':
				return do_extract(optarg, 0);
				break;				
			case 'F':
				return do_extract(optarg, 1);
				break;				
			case 'X':
				return do_check(optarg);
				break;			
			case 'h':
			default:
				usage();
				return 1;
		}
	}	
#endif
	return 0;
}

